import os

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.readers.feishu_docs import FeishuDocsReader
from llama_index.readers.file import PyMuPDFReader

if __name__ == "__main__":
    os.environ["OPENAI_API_KEY"] = os.environ["OPENAI_API_KEY_ZHIHU"]
    os.environ["OPENAI_API_BASE"] = os.environ["OPENAI_API_BASE_ZHIHU"]
    os.environ["OPENAI_BASE_URL"] = os.environ["OPENAI_API_BASE_ZHIHU"]

    import json
    from pydantic.v1 import BaseModel
    def show_json(data):
        """用于展示json数据"""
        if isinstance(data, str):
            obj = json.loads(data)
            print(json.dumps(obj, indent=4))
        elif isinstance(data, dict) or isinstance(data, list):
            print(json.dumps(data, indent=4))
        elif issubclass(type(data), BaseModel):
            print(json.dumps(data.dict(), indent=4, ensure_ascii=False))


    def show_list_obj(data):
        """用于展示一组对象"""
        if isinstance(data, list):
            for item in data:
                show_json(item)
        else:
            raise ValueError("Input is not a list")


    # 见说明文档
    app_id = "cli_a6f1c0fa1fd9d00b"
    app_secret = "dMXCTy8DGaty2xn8I858ZbFDFvcqgiep"

    # https://agiclass.feishu.cn/docx/FULadzkWmovlfkxSgLPcE4oWnPf
    # 链接最后的 "FULadzkWmovlfkxSgLPcE4oWnPf" 为文档 ID
    doc_ids = ["FULadzkWmovlfkxSgLPcE4oWnPf"]

    # 定义飞书文档加载器
    loader = FeishuDocsReader(app_id, app_secret)

    # 加载文档
    documents = loader.load_data(document_ids=doc_ids)

    show_json(documents[0])
    print("==============Document==============\n")
    # 显示前1000字符
    print(documents[0].text[:1000])